# -*- coding: utf-8 -*-
# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""The parallel execution of a BrainPy func on multiple CPU cores.

Specifically, these batch running functions include:

- ``cpu_ordered_parallel``: Performs a parallel ordered map.
- ``cpu_unordered_parallel``: Performs a parallel unordered map.
"""

import sys
from collections.abc import Sized
from typing import (Any, Callable, Generator, Iterable, List,
                    Union, Optional, Sequence, Dict)

from tqdm.auto import tqdm

from brainpy._errors import PackageMissingError

try:
    from pathos.helpers import cpu_count  # noqa
    from pathos.multiprocessing import ProcessPool  # noqa
    import multiprocess.context as ctx  # noqa

    ctx._force_start_method('spawn')
except ModuleNotFoundError:
    cpu_count = None
    ProcessPool = None

__all__ = [
    'cpu_ordered_parallel',
    'cpu_unordered_parallel',
]


def _parallel(
    ordered: bool,
    function: Callable,
    arguments: Union[Sequence[Iterable], Dict[str, Iterable]],
    num_process: Union[int, float] = None,
    num_task: int = None,
    **tqdm_kwargs: Any
) -> Generator:
    """Perform a parallel map with a progress bar.

    Parameters::

    ordered: bool
      True for an ordered map, false for an unordered map.
    function: callable, function
      The function to apply to each element of the given Iterables.
    arguments: sequence of Iterable, dict
      One or more Iterables containing the data to be mapped.
    num_process: int, float
      Number of threads used for parallel running. If `int`, it is
      the number of threads to be used; if `float`, it is the fraction
      of total threads to be used for running.
    num_task: int
      The total number of tasks in this parallel running.
    tqdm_kwargs: Any
      The setting for the progress bar.

    Returns::

    results: Iterable
        A generator which will apply the function to each element of the given Iterables
        in parallel in order with a progress bar.
    """
    if sys.platform == 'win32' and sys.version_info.minor >= 11:
        raise NotImplementedError('Multiprocessing is not available in Python >=3.11 on Windows. '
                                  'Please use Linux or MacOS, or Windows with Python <= 3.10.')

    if ProcessPool is None or cpu_count is None:
        raise PackageMissingError(
            '''
          Please install "pathos" package first. 
          
          >>>  pip install pathos
            '''
        )

    # Determine num_process
    if num_process is None:
        num_process = cpu_count()
    elif isinstance(num_process, int):
        pass
    elif isinstance(num_process, float):
        num_process = int(round(num_process * cpu_count()))
    else:
        raise ValueError('"num_process" must be an int or a float.')

    # arguments
    if isinstance(arguments, dict):
        keys = list(arguments.keys())
        arguments = list(arguments.values())
        run_f = lambda *args: function(**{key: arg for key, arg in zip(keys, args)})
    else:
        if not isinstance(arguments, (tuple, list)):
            raise TypeError('"arguments" must be a sequence of Iterable or a dict of Iterable. '
                            f'But we got {type(arguments)}')
        run_f = function

    # Determine length of tqdm
    lengths = [len(iterable) for iterable in arguments if isinstance(iterable, Sized)]
    num_task = num_task or (min(lengths) if lengths else None)

    # Create parallel generator
    pool = ProcessPool(nodes=num_process)
    if ordered:
        map_func = pool.imap
    else:
        map_func = pool.uimap

    # Choose tqdm variant
    for item in tqdm(map_func(run_f, *arguments), total=num_task, **tqdm_kwargs):
        yield item

    pool.clear()


def cpu_ordered_parallel(
    func: Callable,
    arguments: Union[Sequence[Iterable], Dict[str, Iterable]],
    num_process: Optional[Union[int, float]] = None,
    num_task: Optional[int] = None,
    **tqdm_kwargs: Any
) -> List[Any]:
    """Performs a parallel ordered map with a progress bar.

    Examples::

    >>> import brainpy as bp
    >>> import brainpy.math as bm
    >>> import numpy as np
    >>>
    >>> def simulate(inp):
    >>>   inp = bm.as_jax(inp)
    >>>   hh = bp.dyn.HH(1)
    >>>   runner = bp.DSRunner(hh, inputs=['input', inp],
    >>>                        monitors=['V', 'spike'],
    >>>                        progress_bar=False)
    >>>   runner.run(100)
    >>>   bm.clear_buffer_memory()  # clear all cached data and functions
    >>>   return runner.mon.spike.sum()
    >>>
    >>> if __name__ == '__main__':  # This is important!
    >>>   results = bp.running.cpu_unordered_parallel(simulate, [np.arange(1, 10, 100)], num_process=10)
    >>>   print(results)

    Parameters::

    func: callable, function
      The function to apply to each element of the given Iterables.
    arguments: sequence of Iterable, dict
      One or more Iterables containing the data to be mapped.
    num_process: int, float
      Number of threads used for parallel running. If `int`, it is
      the number of threads to be used; if `float`, it is the fraction
      of total threads to be used for running.
    num_task: int
      The total number of tasks in this parallel running.
    tqdm_kwargs: Any
      The setting for the progress bar.

    Returns::

    results: list
      A list which will apply the function to each element of the given tasks.
    """
    generator = _parallel(True,
                          func,
                          arguments,
                          num_process=num_process,
                          num_task=num_task,
                          **tqdm_kwargs)
    return list(generator)


def cpu_unordered_parallel(
    func: Callable,
    arguments: Union[Sequence[Iterable], Dict[str, Iterable]],
    num_process: Optional[Union[int, float]] = None,
    num_task: Optional[int] = None,
    **tqdm_kwargs: Any
) -> List[Any]:
    """Performs a parallel unordered map with a progress bar.

    Examples::

    >>> import brainpy as bp
    >>> import brainpy.math as bm
    >>> import numpy as np
    >>>
    >>> def simulate(inp):
    >>>   inp = bm.as_jax(inp)
    >>>   hh = bp.dyn.HH(1)
    >>>   runner = bp.DSRunner(hh, inputs=['input', inp],
    >>>                        monitors=['V', 'spike'],
    >>>                        progress_bar=False)
    >>>   runner.run(100)
    >>>   bm.clear_buffer_memory()  # clear all cached data and functions
    >>>   return runner.mon.spike.sum()
    >>>
    >>> if __name__ == '__main__':  # This is important!
    >>>   results = bp.running.cpu_unordered_parallel(simulate, [np.arange(1, 10, 100)], num_process=10)
    >>>   print(results)

    Parameters::

    func: callable, function
      The function to apply to each element of the given Iterables.
    arguments: sequence of Iterable, dict
      One or more Iterables containing the data to be mapped.
    num_process: int, float
      Number of threads used for parallel running. If `int`, it is
      the number of threads to be used; if `float`, it is the fraction
      of total threads to be used for running.
    num_task: int
      The total number of tasks in this parallel running.
    tqdm_kwargs: Any
      The setting for the progress bar.

    Returns::

    results: list
      A list which will apply the function to each element of the given tasks.
    """
    generator = _parallel(False,
                          func,
                          arguments,
                          num_process=num_process,
                          num_task=num_task,
                          **tqdm_kwargs)
    return list(generator)
